/*
Copyright 2012 Martin Sliwka

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

#include "stdafx.h"
#include "AsyncNamedPipe.h"

#define BUFFER_SIZE	(1024*1024)

inline std::string TrimStr(const std::string& String, const std::string& TrimCharacters = " \t\r\n")
{
	int p2 = String.find_last_not_of(TrimCharacters);
	if (p2 == std::string::npos) return std::string();
	int p1 = String.find_first_not_of(TrimCharacters);
	if (p1 == std::string::npos) p1 = 0;
	return String.substr(p1, (p2-p1)+1);
}

std::string GetMessage( HMODULE hModule, DWORD dwLastError )
{
	std::string Result;
	LPSTR MessageBuffer = NULL;
	DWORD dwBufferLength = 0;

	DWORD dwFormatFlags = FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_IGNORE_INSERTS | FORMAT_MESSAGE_FROM_SYSTEM;

	if (hModule != NULL)
		dwFormatFlags |= FORMAT_MESSAGE_FROM_HMODULE;

	if (dwBufferLength = FormatMessageA(dwFormatFlags, hModule, dwLastError, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPSTR)&MessageBuffer, 0, NULL))
	{
		Result = MessageBuffer;
		LocalFree(MessageBuffer);
	}

	Result = TrimStr(Result);

	return Result;
}

AsyncNamedPipeT::AsyncNamedPipeT()
{
	mPipe = NULL;
	mThread = NULL;
	mConnected = NULL;
}

AsyncNamedPipeT::~AsyncNamedPipeT()
{
	SAFE_CLOSEHANDLE(mConnected);
	SAFE_CLOSEHANDLE(mThread);
	SAFE_CLOSEHANDLE(mPipe);
}

HRESULT AsyncNamedPipeT::Create( LPCSTR pName, bool FirstInstance )
{
	mFirstInstance = FirstInstance;

	LOG("[%d] Creating pipe '%s'...\n", !mFirstInstance, pName);

	DWORD LastError = NO_ERROR;

	DWORD dwOpenMode = PIPE_ACCESS_DUPLEX;
	if (mFirstInstance)
		dwOpenMode |= FILE_FLAG_FIRST_PIPE_INSTANCE;

	DWORD dwPipeMode = PIPE_TYPE_BYTE | PIPE_READMODE_BYTE;

	mPipe = ::CreateNamedPipeA(pName, dwOpenMode, dwPipeMode, 2, 0, BUFFER_SIZE, BUFFER_SIZE, nullptr);
	if (mPipe == INVALID_HANDLE_VALUE)
		mPipe = NULL;
	if (!mPipe)
		LastError = HRESULT_FROM_WIN32(GetLastError());

	LOG("[%d] %s (%d)\n", !mFirstInstance, GetMessage(NULL, LastError).c_str(), LastError);

	mConnected = ::CreateEvent(nullptr, TRUE, FALSE, nullptr);

	return HRESULT_FROM_WIN32(LastError);
}

HRESULT AsyncNamedPipeT::StartTransfer( AsyncNamedPipePtr Writer )
{
	DWORD LastError = NO_ERROR;

	mOtherPipe = Writer;

	LOG("[%d] Creating thread...\n", !mFirstInstance);

	mThread = reinterpret_cast<HANDLE>(_beginthreadex(NULL, 0, ThreadProcWrapper, this, CREATE_SUSPENDED, NULL));
	if (!mThread)
		LastError = ERROR_TOO_MANY_TCBS;

	LOG("[%d] %s (%d)\n", !mFirstInstance, GetMessage(NULL, LastError).c_str(), LastError);

	if (mThread)
		::ResumeThread(mThread);

	return S_OK;

}
HRESULT AsyncNamedPipeT::WaitForCompletion( DWORD Timeout )
{
	DWORD WaitResult = ::WaitForSingleObject(mThread, Timeout);

	if (WaitResult == WAIT_TIMEOUT)
		return S_FALSE;

	if (WaitResult == WAIT_OBJECT_0)
		return S_OK;

	return E_FAIL;
}

void AsyncNamedPipeT::Destroy()
{
	SAFE_CLOSEHANDLE(mPipe);
	SAFE_CLOSEHANDLE(mThread);
}

unsigned int WINAPI AsyncNamedPipeT::ThreadProcWrapper( LPVOID lpParameter )
{
	return reinterpret_cast<AsyncNamedPipeT*>(lpParameter)->ThreadProc();
}

unsigned int AsyncNamedPipeT::ThreadProc()
{
	HRESULT Result = S_OK;

	bool Connected = false;

	const DWORD BufferSize = BUFFER_SIZE;
	LPBYTE pBuffer = new BYTE[BufferSize];
	DWORD BytesAvailable = 0;
	DWORD BytesRead = 0;
	DWORD BytesWritten = 0;

	LOG("[%d] Connecting...\n", !mFirstInstance);

	if (SUCCEEDED(Result))
		Result = Connect();

	if (SUCCEEDED(Result))
		LOG("[%d] Waiting for data...\n", !mFirstInstance);

	while (SUCCEEDED(Result))
	{
		bool PipeBroken = false;

		Result = Read(pBuffer, BufferSize, BytesRead);
		
		if (Result == S_FALSE)
			PipeBroken = true;

		if (BytesRead > 0)
		{
			Result = mOtherPipe->Write(pBuffer, BytesRead, BytesWritten);

			if (Result == S_FALSE)
				PipeBroken = true;
		}
		else
		{
			Sleep(10);
		}

		if (PipeBroken)
			Result = E_ABORT;
	}

	Disconnect();
	mOtherPipe->Disconnect();

	return 0;
}

//////////////////////////////////////////////////////////////////////////

HRESULT AsyncNamedPipeT::Connect()
{
	DWORD LastError = NO_ERROR;

	if (!::ConnectNamedPipe(mPipe, nullptr))
		LastError = GetLastError();

	LOG("[%d] %s (%d)\n", !mFirstInstance, GetMessage(NULL, LastError).c_str(), LastError);

	switch (LastError)
	{
	case NO_ERROR:
	case ERROR_PIPE_CONNECTED:	// client is still connected
	case ERROR_NO_DATA:			// client was fast enought to disconnect before we "connect"
	case ERROR_PIPE_LISTENING:	// in case of "nowait" mode
		::SetEvent(mConnected);
		return S_OK;
	}

	assert(false);
	return HRESULT_FROM_WIN32(LastError);
}

HRESULT AsyncNamedPipeT::Peek( DWORD& BytesAvailable )
{
	DWORD LastError = NO_ERROR;

	BytesAvailable = 0;
	if (!::PeekNamedPipe(mPipe, nullptr, 0, nullptr, &BytesAvailable, nullptr))
		LastError = GetLastError();

	LOG("[%d] Available %d bytes, %s (%d)\n", !mFirstInstance, BytesAvailable, GetMessage(NULL, LastError).c_str(), LastError);

	switch (LastError)
	{
	case ERROR_NO_DATA:			// client disconnected
	case ERROR_BROKEN_PIPE:
		return S_FALSE;

	case NO_ERROR:
		return S_OK;
	}

	assert(false);
	return HRESULT_FROM_WIN32(LastError);
}

HRESULT AsyncNamedPipeT::Read( LPBYTE pBuffer, DWORD BytesToRead, DWORD& BytesRead )
{
	DWORD LastError = NO_ERROR;

	DWORD BytesAvailable = 0;
	if (!::PeekNamedPipe(mPipe, nullptr, 0, nullptr, &BytesAvailable, nullptr))
		LastError = GetLastError();

	BytesRead = 0;
	if (LastError == NO_ERROR && BytesAvailable > 0)
	{
		if (!::ReadFile(mPipe, pBuffer, BytesToRead, &BytesRead, nullptr))
			LastError = GetLastError();

		LOG("[%d] Read %d bytes, %s (%d)\n", !mFirstInstance, BytesRead, GetMessage(NULL, LastError).c_str(), LastError);
	}

	switch (LastError)
	{
	case ERROR_NO_DATA:			// client disconnected
	case ERROR_BROKEN_PIPE:
	case ERROR_BAD_PIPE:
		return S_FALSE;

	case NO_ERROR:
		return S_OK;
	}

	assert(false);
	return HRESULT_FROM_WIN32(LastError);
}

HRESULT AsyncNamedPipeT::Write( LPBYTE pBuffer, DWORD BytesToWrite, DWORD& BytesWritten )
{
	DWORD LastError = NO_ERROR;

	BytesWritten = 0;
	if (!::WriteFile(mPipe, pBuffer, BytesToWrite, &BytesWritten, nullptr))
		LastError = GetLastError();

	LOG("[%d] Written %d bytes, %s (%d)\n", !mFirstInstance, BytesWritten, GetMessage(NULL, LastError).c_str(), LastError);

	switch (LastError)
	{
	case ERROR_NO_DATA:			// client disconnected
	case ERROR_BROKEN_PIPE:
	case ERROR_BAD_PIPE:
		return S_FALSE;

	case NO_ERROR:
		return S_OK;
	}

	assert(false);
	return HRESULT_FROM_WIN32(LastError);
}

void AsyncNamedPipeT::Disconnect()
{
	::FlushFileBuffers(mPipe);
	::DisconnectNamedPipe(mPipe);
}

